import torch
from dataset import test_loader, train_loader
from model import Model

if __name__ == '__main__':
    model = Model()
    model.load_state_dict(torch.load('output/mnist1.pth'))

    errors: list = []
    for image, target_label in train_loader:
        output = torch.flatten(model(image))

        for i, n in enumerate(output):
            if n == max(output):
                output_label = i

        if output_label != int(target_label):
            errors.append((output_label, int(target_label)))
    print(len(errors))